import re
import os
import logging
import requests
from langchain_community.llms import QianfanLLMEndpoint
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.agents import create_react_agent, Tool
from langchain.prompts import PromptTemplate
from langchain.chains import SequentialChain

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 设置凭证
os.environ['QIANFAN_AK'] = 'Kdi2l1hCbBA87gqiS7dsOKu7'
os.environ['QIANFAN_SK'] = 'FbJg0n1suYnzaVopylXGKy4TMtYphTpv'


class SearchRun:
    def __init__(self, engine_name, query_url):
        self.engine_name = engine_name
        self.query_url = query_url

    def run(self, query):
        try:
            url = self.query_url.format(query=query)
            response = requests.get(url)
            if response.status_code == 200:
                return response.text
            else:
                return f"{self.engine_name}搜索请求失败。"
        except Exception as e:
            logging.error(f"{self.engine_name}搜索时出错: {e}")
            return f"调用{self.engine_name}搜索过程中发生错误。"


def initialize_agents_and_tools():
    chat_model_1 = QianfanLLMEndpoint(model="Qianfan-Chinese-Llama-2-7B")
    search_tool = DuckDuckGoSearchRun()
    baidu_tool = SearchRun(engine_name="Baidu", query_url="https://www.baidu.com/s?wd={query}")
    return chat_model_1, search_tool, baidu_tool


# 初始化助手和搜索工具
chat_model_1, search_tool, baidu_tool = initialize_agents_and_tools()

# 定义可用工具
tools = [
    Tool(name="DuckDuckGo Search", func=search_tool.run, description="用于实时搜索信息"),
    Tool(name="Baidu Search", func=baidu_tool.run, description="用于百度搜索信息")
]

# 创建提示模板
prompt = PromptTemplate(
    input_variables=["input", "tools", "agent_scratchpad", "tool_names"],
    template="""\
你是一个智能助手。请根据用户的问题给出直接的答案，避免冗长的解释和条件说明。
用户的问题是：{input}

助手记录本：{agent_scratchpad}  
可用工具：{tools}  
工具名称：{tool_names}  

请直接回答用户的问题。请使用以下格式返回：
- 直接回答：<回答内容>
- 如果需要搜索：search("你的查询")
- 如果没有答案，直接说明：<没有找到>
"""
)

# 创建代理
agent = create_react_agent(llm=chat_model_1, tools=tools, prompt=prompt)

# 定义生成题目的模板
prompt_generate = PromptTemplate(
    template="请你告知我病情如何治疗{question}",
    input_variables=["question"]
)

# 修改后的生成链，使用 invoke() 方法代替 run()
generate_chain = prompt_generate | chat_model_1

# 创建COT推理模板
cot_prompt = PromptTemplate(
    input_variables=["question"],
    template="""\
你是一个有逻辑推理能力的助手。请分步骤回答以下问题：
问题：{question}

第一步：分析问题的主要要素。
第二步：提出假设并验证。
第三步：总结推理结果并提供答案。
"""
)

# 修改后的COT推理链
cot_chain = cot_prompt | chat_model_1


# 定义上下文管理类
class QueryHandler:
    def __init__(self):
        self.context = {}

    def handle_query(self, query, agent_scratchpad):
        # 判断是否是医学相关问题
        if self.is_medical_query(query):
            response = self.handle_medical_query(query)
        else:
            if query in self.context:
                return self.context[query]
            else:
                if "复杂" in query:
                    response = self.cot_reasoning(query)
                else:
                    response = self.process_query_with_model(query, chat_model_1, agent_scratchpad)
                self.context[query] = response
        return response

    def is_medical_query(self, query):
        # 简单的医学关键词检测
        medical_keywords = ["头疼", "发烧", "感冒", "疼痛", "咳嗽"]
        return any(keyword in query for keyword in medical_keywords)

    def handle_medical_query(self, query):
        # 处理医学相关的查询，改为使用 invoke() 方法
        try:
            response = generate_chain.invoke({"question": query})
            response = response.strip()
            logging.info(f"医学查询: {query}")
            logging.info(f"原始响应: {response}")
            return self.parse_response(response)
        except Exception as e:
            logging.error(f"处理医学查询时出错: {e}")
            return "处理医学请求时发生错误。"

    def process_query_with_model(self, query, model, agent_scratchpad):
        try:
            response = generate_chain.invoke({"question": query})
            response = response.strip()
            logging.info(f"查询: {query}")
            logging.info(f"原始响应: {response}")
            return self.parse_response(response)
        except Exception as e:
            logging.error(f"处理查询时出错: {e}")
            return "处理请求时发生错误。"

    def parse_response(self, response):
        logging.info(f"解析响应: {response}")  # 记录原始响应
        direct_answer_match = re.search(r'-\s*直接回答：\s*(.+?)(?=-|如果需要搜索：|如果没有答案，直接说明：|$)', response,
                                        re.DOTALL)

        if direct_answer_match:
            return direct_answer_match.group(1).strip()

        search_match = re.search(r'如果需要搜索：search\("(.+?)"\)', response)
        if search_match:
            query = search_match.group(1)
            return f"请使用搜索功能：{query}"

        # 记录没有匹配到的情况
        logging.warning("没有找到匹配的回答格式。")
        return "无法解析回答。"


if __name__ == "__main__":
    handler = QueryHandler()
    question = handler.handle_query("头疼一般是如何引发的？", "")
    print(question)


